{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t0CF6Gvkt_Cw"
   },
   "source": [
    "# MuJoCo basics\n",
    "\n",
    "We begin by defining and loading a simple model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mujoco\n",
    "import mediapy as media\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import time\n",
    "import itertools\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1kOs1wTc7uCZ"
   },
   "source": [
    "# Advanced rendering\n",
    "\n",
    "Like joint visualization, additional rendering options are exposed as parameters to the `render` method.\n",
    "\n",
    "Let's bring back our first model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "mTDgsk2xcgwH"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table class=\"show_images\" style=\"border-spacing:0px;\"><tr><td style=\"padding:1px;\"><img width=\"320\" height=\"240\" style=\"image-rendering:auto; object-fit:cover;\" src=\"\"/></td></tr></table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "xml = \"\"\"\n",
    "<mujoco>\n",
    "  <worldbody>\n",
    "    <light name=\"top\" pos=\"0 0 1\"/>\n",
    "    <body name=\"box_and_sphere\" euler=\"0 0 -30\">\n",
    "      <joint name=\"swing\" type=\"hinge\" axis=\"1 -1 0\" pos=\"-.2 -.2 -.2\"/>\n",
    "      <geom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/>\n",
    "      <geom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/>\n",
    "    </body>\n",
    "  </worldbody>\n",
    "</mujoco>\n",
    "\"\"\"\n",
    "model = mujoco.MjModel.from_xml_string(xml)\n",
    "renderer = mujoco.Renderer(model)\n",
    "data = mujoco.MjData(model)\n",
    "\n",
    "mujoco.mj_forward(model, data)\n",
    "renderer.update_scene(data)\n",
    "media.show_image(renderer.render())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "VePXamL_6XUc"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table class=\"show_images\" style=\"border-spacing:0px;\"><tr><td style=\"padding:1px;\"><img width=\"320\" height=\"240\" style=\"image-rendering:auto; object-fit:cover;\" src=\"\"/></td></tr></table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#@title Enable transparency and frame visualization\n",
    "\n",
    "# enable joint visualization option:\n",
    "scene_option = mujoco.MjvOption()\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = True\n",
    "scene_option.frame = mujoco.mjtFrame.mjFRAME_GEOM\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = True\n",
    "renderer.update_scene(data, scene_option=scene_option)\n",
    "frame = renderer.render()\n",
    "media.show_image(frame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PVcpcvww9lZ8"
   },
   "outputs": [],
   "source": [
    "#@title Depth rendering\n",
    "\n",
    "# update renderer to render depth\n",
    "renderer.enable_depth_rendering()\n",
    "\n",
    "# reset the scene\n",
    "renderer.update_scene(data)\n",
    "\n",
    "# depth is a float array, in meters.\n",
    "depth = renderer.render()\n",
    "\n",
    "# Shift nearest values to the origin.\n",
    "depth -= depth.min()\n",
    "# Scale by 2 mean distances of near rays.\n",
    "depth /= 2*depth[depth <= 1].mean()\n",
    "# Scale to [0, 255]\n",
    "pixels = 255*np.clip(depth, 0, 1)\n",
    "\n",
    "media.show_image(pixels.astype(np.uint8))\n",
    "\n",
    "renderer.disable_depth_rendering()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PNwiIrgpx7T8"
   },
   "outputs": [],
   "source": [
    "#@title Segmentation rendering\n",
    "\n",
    "# update renderer to render segmentation\n",
    "renderer.enable_segmentation_rendering()\n",
    "\n",
    "# reset the scene\n",
    "renderer.update_scene(data)\n",
    "\n",
    "seg = renderer.render()\n",
    "\n",
    "# Display the contents of the first channel, which contains object\n",
    "# IDs. The second channel, seg[:, :, 1], contains object types.\n",
    "geom_ids = seg[:, :, 0]\n",
    "# Infinity is mapped to -1\n",
    "geom_ids = geom_ids.astype(np.float64) + 1\n",
    "# Scale to [0, 1]\n",
    "geom_ids = geom_ids / geom_ids.max()\n",
    "pixels = 255*geom_ids\n",
    "media.show_image(pixels.astype(np.uint8))\n",
    "\n",
    "renderer.disable_segmentation_rendering()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wo72mo0mGIXr"
   },
   "source": [
    "## The camera matrix\n",
    "\n",
    "For a description of the camera matrix see the article [Camera matrix](https://en.wikipedia.org/wiki/Camera_matrix) on Wikipedia."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sDYwClpxaxab"
   },
   "outputs": [],
   "source": [
    "def compute_camera_matrix(renderer, data):\n",
    "  \"\"\"Returns the 3x4 camera matrix.\"\"\"\n",
    "  # If the camera is a 'free' camera, we get its position and orientation\n",
    "  # from the scene data structure. It is a stereo camera, so we average over\n",
    "  # the left and right channels. Note: we call `self.update()` in order to\n",
    "  # ensure that the contents of `scene.camera` are correct.\n",
    "  renderer.update_scene(data)\n",
    "  pos = np.mean([camera.pos for camera in renderer.scene.camera], axis=0)\n",
    "  z = -np.mean([camera.forward for camera in renderer.scene.camera], axis=0)\n",
    "  y = np.mean([camera.up for camera in renderer.scene.camera], axis=0)\n",
    "  rot = np.vstack((np.cross(y, z), y, z))\n",
    "  fov = model.vis.global_.fovy\n",
    "\n",
    "  # Translation matrix (4x4).\n",
    "  translation = np.eye(4)\n",
    "  translation[0:3, 3] = -pos\n",
    "\n",
    "  # Rotation matrix (4x4).\n",
    "  rotation = np.eye(4)\n",
    "  rotation[0:3, 0:3] = rot\n",
    "\n",
    "  # Focal transformation matrix (3x4).\n",
    "  focal_scaling = (1./np.tan(np.deg2rad(fov)/2)) * renderer.height / 2.0\n",
    "  focal = np.diag([-focal_scaling, focal_scaling, 1.0, 0])[0:3, :]\n",
    "\n",
    "  # Image matrix (3x3).\n",
    "  image = np.eye(3)\n",
    "  image[0, 2] = (renderer.width - 1) / 2.0\n",
    "  image[1, 2] = (renderer.height - 1) / 2.0\n",
    "  return image @ focal @ rotation @ translation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Bs89vS0wLoU0"
   },
   "source": [
    "Let's use the camera matrix to project from world to camera coordinates:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "My0N4_7PDJ_q"
   },
   "outputs": [],
   "source": [
    "# reset the scene\n",
    "renderer.update_scene(data)\n",
    "\n",
    "\n",
    "# Get the world coordinates of the box corners\n",
    "box_pos = data.geom_xpos[model.geom('red_box').id]\n",
    "box_mat = data.geom_xmat[model.geom('red_box').id].reshape(3, 3)\n",
    "box_size = model.geom_size[model.geom('red_box').id]\n",
    "offsets = np.array([-1, 1]) * box_size[:, None]\n",
    "xyz_local = np.stack(list(itertools.product(*offsets))).T\n",
    "xyz_global = box_pos[:, None] + box_mat @ xyz_local\n",
    "\n",
    "# Camera matrices multiply homogenous [x, y, z, 1] vectors.\n",
    "corners_homogeneous = np.ones((4, xyz_global.shape[1]), dtype=float)\n",
    "corners_homogeneous[:3, :] = xyz_global\n",
    "\n",
    "# Get the camera matrix.\n",
    "m = compute_camera_matrix(renderer, data)\n",
    "\n",
    "# Project world coordinates into pixel space. See:\n",
    "# https://en.wikipedia.org/wiki/3D_projection#Mathematical_formula\n",
    "xs, ys, s = m @ corners_homogeneous\n",
    "# x and y are in the pixel coordinate system.\n",
    "x = xs / s\n",
    "y = ys / s\n",
    "\n",
    "# Render the camera view and overlay the projected corner coordinates.\n",
    "pixels = renderer.render()\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "ax.imshow(pixels)\n",
    "ax.plot(x, y, '+', c='w')\n",
    "ax.set_axis_off()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AGm5-e0sHEAF"
   },
   "source": [
    "## Modifying the scene\n",
    "\n",
    "Let's add some arbitrary geometry to the `mjvScene`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z6NDYJ8IOVt7"
   },
   "outputs": [],
   "source": [
    "def get_geom_speed(model, data, geom_name):\n",
    "  \"\"\"Returns the speed of a geom.\"\"\"\n",
    "  geom_vel = np.zeros(6)\n",
    "  geom_type = mujoco.mjtObj.mjOBJ_GEOM\n",
    "  geom_id = data.geom(geom_name).id\n",
    "  mujoco.mj_objectVelocity(model, data, geom_type, geom_id, geom_vel, 0)\n",
    "  return np.linalg.norm(geom_vel)\n",
    "\n",
    "def add_visual_capsule(scene, point1, point2, radius, rgba):\n",
    "  \"\"\"Adds one capsule to an mjvScene.\"\"\"\n",
    "  if scene.ngeom >= scene.maxgeom:\n",
    "    return\n",
    "  scene.ngeom += 1  # increment ngeom\n",
    "  # initialise a new capsule, add it to the scene using mjv_makeConnector\n",
    "  mujoco.mjv_initGeom(scene.geoms[scene.ngeom-1],\n",
    "                      mujoco.mjtGeom.mjGEOM_CAPSULE, np.zeros(3),\n",
    "                      np.zeros(3), np.zeros(9), rgba.astype(np.float32))\n",
    "  mujoco.mjv_makeConnector(scene.geoms[scene.ngeom-1],\n",
    "                           mujoco.mjtGeom.mjGEOM_CAPSULE, radius,\n",
    "                           point1[0], point1[1], point1[2],\n",
    "                           point2[0], point2[1], point2[2])\n",
    "\n",
    " # traces of time, position and speed\n",
    "times = []\n",
    "positions = []\n",
    "speeds = []\n",
    "offset = model.jnt_axis[0]/8  # offset along the joint axis\n",
    "\n",
    "def modify_scene(scn):\n",
    "  \"\"\"Draw position trace, speed modifies width and colors.\"\"\"\n",
    "  if len(positions) > 1:\n",
    "    for i in range(len(positions)-1):\n",
    "      rgba=np.array((np.clip(speeds[i]/10, 0, 1),\n",
    "                     np.clip(1-speeds[i]/10, 0, 1),\n",
    "                     .5, 1.))\n",
    "      radius=.003*(1+speeds[i])\n",
    "      point1 = positions[i] + offset*times[i]\n",
    "      point2 = positions[i+1] + offset*times[i+1]\n",
    "      add_visual_capsule(scn, point1, point2, radius, rgba)\n",
    "\n",
    "duration = 6    # (seconds)\n",
    "framerate = 30  # (Hz)\n",
    "\n",
    "# Simulate and display video.\n",
    "frames = []\n",
    "\n",
    "# Reset state and time.\n",
    "mujoco.mj_resetData(model, data)\n",
    "mujoco.mj_forward(model, data)\n",
    "\n",
    "while data.time < duration:\n",
    "  # append data to the traces\n",
    "  positions.append(data.geom_xpos[data.geom(\"green_sphere\").id].copy())\n",
    "  times.append(data.time)\n",
    "  speeds.append(get_geom_speed(model, data, \"green_sphere\"))\n",
    "  mujoco.mj_step(model, data)\n",
    "  if len(frames) < data.time * framerate:\n",
    "    renderer.update_scene(data)\n",
    "    modify_scene(renderer.scene)\n",
    "    pixels = renderer.render()\n",
    "    frames.append(pixels)\n",
    "media.show_video(frames, fps=framerate)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "-re3Szx-1Ias"
   ],
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "gpuClass": "premium",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
